#include <crtdbg.h>

#include "tdp2_threading.hpp"

namespace TDP
{
	namespace Threading
	{
		void SimpleThreadDispatcher::RunSimpleJobProc::Start() const
		{
			for(void *work=_job->GetWork();work != NULL;work = _job->GetWork())
				_job->Work(work);
			AtomicIncrement(_availableThreadCounter);
		}

		SimpleThreadDispatcher::SimpleThreadDispatcher()
		{
			_numTotalThreads = CountProcessors();
			Threading::NonAtomicWrite(&_numAvailableThreads, _numTotalThreads);
		}

		SimpleThreadDispatcher::~SimpleThreadDispatcher()
		{
		}

		void SimpleThreadDispatcher::DigestJob(SimpleJob *job, bool synchronous, bool workUsingCurrent)
		{
			ISystemThread **threads = NULL;
			RunSimpleJobProc **procs = NULL;
			int numThreadsAllocated = 0;

			if(!synchronous)
			{
				// Spawn work threads
				threads = new ISystemThread *[_numTotalThreads];
				procs = new RunSimpleJobProc *[_numTotalThreads];
				while(AtomicDecrement(&_numAvailableThreads) >= 0)
				{
					procs[numThreadsAllocated] = new RunSimpleJobProc(&_numAvailableThreads, job);
					threads[numThreadsAllocated] = CreateSystemThread(procs[numThreadsAllocated]);
					numThreadsAllocated++;
					if(numThreadsAllocated >= _numTotalThreads)
					{
						// Work threads completed while this was running, causing numAvailableThreads to increase
						// Decrement the counter to keep the count in sync and get out
						AtomicDecrement(&_numAvailableThreads);
						break;
					}
				}
				AtomicIncrement(&_numAvailableThreads);
			}
			if(synchronous || workUsingCurrent || !numThreadsAllocated)
			{
				// Work in this thread
				for(void *work=job->GetWork();work != NULL;work = job->GetWork())
				{
					job->Work(work);
					job->Twiddle();
				}
			}

			// Wait for worker threads to come back
			for(int i=0;i<numThreadsAllocated;i++)
			{
				while(!threads[i]->WaitFor(1000))
					job->Twiddle();
				TDP::Threading::DestroySystemThread(threads[i]);
				delete procs[i];
			}

			if(threads) delete[] threads;
			if(procs) delete[] procs;

			// Twiddle one last time to potentially update progress reporting
			job->Twiddle();
		}

		class MultilevelThreadWorkProc : public ISystemThreadProc
		{
			MultilevelThreadDispatcher *_owner;
			volatile AtomicInt *_startCounter;

		public:
			explicit MultilevelThreadWorkProc(MultilevelThreadDispatcher *owner, volatile AtomicInt *startCounter)
			{
				_owner = owner;
				_startCounter = startCounter;
			}

			void Start() const
			{
				MultilevelThreadDispatcher *disp = _owner;
				AtomicDecrement(_startCounter);

				while(disp->IsActive())
				{
					Threading::SleepThread(50);

					MultilevelThreadDispatcher::DispatchableJob *job = disp->AcquireJob();
					if(job)
						job->Digest(false);
				}
			}
		};

		MultilevelThreadDispatcher::MultilevelThreadDispatcher(int numThreads)
		{
			volatile AtomicInt startCounter;
			Threading::NonAtomicWrite(&startCounter, numThreads);
			Threading::NonAtomicWrite(&_active, 1);
			MultilevelThreadWorkProc proc(this, &startCounter);
			_firstJob = _lastJob = NULL;

			_threads = new ISystemThread*[numThreads];
			_numThreads = numThreads;

			for(int i=0;i<numThreads;i++)
				_threads[i] = Threading::CreateSystemThread(&proc);

			// Wait for the threads to spin up
			while(Threading::AtomicRead(&startCounter) != 0)
				Threading::SleepThread(20);
		}

		MultilevelThreadDispatcher::~MultilevelThreadDispatcher()
		{
			Threading::AtomicWrite(&_active, 0);

			// Wait for work threads to wind down
			for(int i=0;i<_numThreads;i++)
			{
				while(_threads[i]->IsActive())
					Threading::SleepThread(20);
				Threading::DestroySystemThread(_threads[i]);
			}
			delete _threads;
		}

		MultilevelThreadDispatcher::DispatchableJob *MultilevelThreadDispatcher::AcquireJob()
		{
			MutexLock<LightMutex> _(&_mutex);

			// Clean up the job list
			while(_firstJob && Threading::NonAtomicRead(&_firstJob->workExhausted) > 0 && Threading::NonAtomicRead(&_firstJob->ownedThreads) == 0 && Threading::NonAtomicRead(&_firstJob->finalize) == 1)
			{
				DispatchableJob *job = _firstJob;
				if(_lastJob == _firstJob)
					_lastJob = _firstJob = NULL;
				else
					_firstJob = _firstJob->nextJob;

				delete job;
			}

			DispatchableJob *job = _firstJob;
			while(job)
			{
				if(Threading::NonAtomicRead(&job->workExhausted) == 0)
				{
					Threading::AtomicIncrement(&job->ownedThreads);
					return job;
				}

				job = job->nextJob;
			}

			return NULL;
		}

		void MultilevelThreadDispatcher::InstallJob(SimpleJob *job, bool twiddle)
		{
			DispatchableJob *djob = new DispatchableJob();
			{
				MutexLock<LightMutex> _(&_mutex);
				djob->job = job;
				djob->nextJob = NULL;
				Threading::NonAtomicWrite(&djob->finalize, 0);
				Threading::NonAtomicWrite(&djob->ownedThreads, 1);
				Threading::NonAtomicWrite(&djob->workExhausted, 0);

				if(_firstJob == NULL)
					_firstJob = djob;
				if(_lastJob != NULL)
					_lastJob->nextJob = djob;
				_lastJob = djob;
				Threading::Synchronize();
			}

			djob->Digest(twiddle);
			while(Threading::NonAtomicRead(&djob->ownedThreads) != 0)
			{
				if(twiddle)
					job->Twiddle();
				Threading::SleepThread(20);
			}
			// Job is completely exhausted, destroy it
			Threading::AtomicWrite(&djob->finalize, 1);
			if(twiddle)
				job->Twiddle();
		}

		bool MultilevelThreadDispatcher::IsActive() const
		{
			return Threading::NonAtomicRead(&this->_active) != 0;
		}
		
		void MultilevelThreadDispatcher::DispatchableJob::Digest(bool twiddle)
		{
			void *work;
			while(true)
			{
				work = this->job->GetWork();
				if(!work)
					break;
				this->job->Work(work);

				if(twiddle)
					this->job->Twiddle();
			}
			Threading::AtomicIncrement(&this->workExhausted);
			Threading::AtomicDecrement(&this->ownedThreads);
		}
	};
};